package com.hjzd.ai.smartpractice.llm;

import org.springframework.http.client.reactive.ClientHttpConnector;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import reactor.core.publisher.Flux;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.netty.http.client.HttpClient;

import java.util.ArrayList;
import java.util.List;

public class Demo {
    public static void main(String[] args) {
        // 创建一个 WebClient 实例
        HttpClient httpClient = HttpClient.create();
        ClientHttpConnector connector = new ReactorClientHttpConnector(httpClient);
        WebClientConfig webClientConfig = new WebClientConfig();
        WebClient webClient = webClientConfig.webClientBuilder()
                .clientConnector(connector)
                .defaultHeader(HttpHeaders.CONTENT_TYPE, "application/json")
                .build();

        // 构造请求的 JSON 字符串
        ChatCompletionMessage chatCompletionMessage0 = new ChatCompletionMessage();
        chatCompletionMessage0.setRole("user");
        chatCompletionMessage0.setContent("你好，给我说个故事");
        ChatCompletionMessage chatCompletionMessage1 = new ChatCompletionMessage();
        chatCompletionMessage1.setRole("assistant");
        chatCompletionMessage1.setContent("你要什么故事");
        ChatCompletionMessage chatCompletionMessage2 = new ChatCompletionMessage();
        chatCompletionMessage2.setRole("user");
        chatCompletionMessage2.setContent("我要恐怖故事");
        List<ChatCompletionMessage> chatCompletionMessages = new ArrayList<>();
        chatCompletionMessages.add(chatCompletionMessage0);
        chatCompletionMessages.add(chatCompletionMessage1);
        chatCompletionMessages.add(chatCompletionMessage2);
        ChatCompletionRequest chatCompletionRequest = new ChatCompletionRequest(chatCompletionMessages, 8192, true,
                "ep-20250205150327-vrtfb", "497e44e4-293f-43b1-96c3-28b6e295c36a");
        // 发起 POST 请求并获取流式返回
        Flux<String> responseStream = webClient.post()
//                .uri("https://ark.cn-beijing.volces.com/api/v3/chat/completions")
//                .uri("http://127.0.0.1:5004/prologue")
                .uri("http://10.253.3.149:5004/chat")
                .header(HttpHeaders.AUTHORIZATION, "Bearer 631999ff-4583-4944-9511-da68d2a58911")
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.TEXT_EVENT_STREAM) // 指定流式响应
                .bodyValue(chatCompletionRequest)
                .retrieve()
                .bodyToFlux(String.class)
                .doOnSubscribe(s -> System.out.println("Waiting for response..."));;

        // 订阅流式响应并打印输出
        responseStream.subscribe(
                data -> System.out.println("Received: " + data),
                error -> System.err.println("Error: " + error),
                () -> System.out.println("Completed")
        );

        // 为了让 demo 程序不会立刻结束，简单等待 15 秒钟
        try {
            Thread.sleep(5000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

}
